import { ref, watch, type Ref } from 'vue-demi'
import { type Fn, type MaybeRefOrGetter, isClient, isWorker, toRef, tryOnScopeDispose, useIntervalFn } from '@vueuse/shared'
import { useEventListener } from '@vueuse/core'

import { isNullAndUnDef } from '@/utils/is'

export type WebSocketStatus = 'OPEN' | 'CONNECTING' | 'CLOSED'

const DEFAULT_PING_MESSAGE = 'ping'

/**
 *
 * @see https://developer.mozilla.org/en-US/docs/Web/API/WebSocket
 */
export interface UseWebSocketOptions {
  onConnected?: (ws: WebSocket) => void
  onDisconnected?: (ws: WebSocket, event: CloseEvent) => void
  onError?: (ws: WebSocket, event: Event) => void
  onMessage?: (ws: WebSocket, event: MessageEvent) => void

  /**
   * Send heartbeat for every x milliseconds passed
   * 每隔 x ms 发送一次心跳
   *
   * @default false
   */
  heartbeat?:
    | boolean
    | {
        /**
         * The message to send
         * 发送的消息
         *
         * @default 'ping'
         */
        message?: string | ArrayBuffer | Blob

        /**
         * Interval, in milliseconds
         * 每隔多少毫秒发送一次心跳
         *
         * @default 1000
         */
        interval?: number

        /**
         * Heartbeat response timeout, in milliseconds
         * 心跳响应超时时间，单位毫秒
         *
         * @default 1000
         */
        pongTimeout?: number
      }

  /**
   * Enabled auto reconnect
   * 是否开启自动重连
   *
   * @default false
   */
  autoReconnect?:
    | boolean
    | {
        /**
         * Maximum retry times.
         * Or you can pass predicate function (which returns true if you want to retry).
         * 最大重试次数
         * 或者传入一个断言函数（布尔表达式），当函数返回 true 时重连
         *
         * @default -1
         */
        retries?: number | (() => boolean)

        /**
         * Delay for reconnect, in milliseconds
         * 重连延迟，单位毫秒
         *
         * @default 1000
         */
        delay?: number

        /**
         * On maximum retry times reached
         * 达到最大重试次数时触发
         */
        onFailed?: Fn
      }

  /**
   * Automatically open a connection
   * 是否自动打开连接
   *
   * @default true
   */
  immediate?: boolean

  /**
   * Automatically close a connection
   * 是否自动关闭连接
   *
   * @default true
   */
  autoClose?: true

  /**
   * List of one or more sub-protocol string
   * 一个或多个子协议字符串
   *
   * @default []
   */
  protocols?: string[]
}

export interface UseWebSocketReturn<T> {
  /**
   * 通过 websocket 接收到的最新数据的引用，
   * 可以监听它以响应传入的消息
   */
  data: Ref<T | null>

  /**
   * 当前 websocket 状态可能的值：
   * 'OPEN', 'CONNECTING', 'CLOSED'
   */
  status: Ref<WebSocketStatus>

  /**
   * 关闭 websocket 连接
   */
  close: WebSocket['close']

  /**
   * 重新打开 websocket 连接。
   * 如果当前连接处于活动状态，则在打开新连接之前将其关闭。
   */
  open: Fn

  /** Sends data through the websocket connection.
   * 通过 websocket 连接发送数据。
   *
   * @param data
   * @param useBuffer 当 socket 连接未打开时，将数据存储到缓冲区并在连接时发送。默认为 true。
   */
  send: (data: string | ArrayBuffer | Blob, useBuffer?: boolean) => boolean

  /**
   * WebSocket 实例的引用。
   */
  ws: Ref<WebSocket | undefined>
}

/**
 * 处理嵌套的选项
 * @param options
 */
function resolveNestedOptions<T>(options: T | true): T {
  if (options === true) {
    return {} as T
  }
  return options
}

/**
 * 响应式的 WebSocket 客户端。
 * @useage
 * @see hooks/docs/useWebsocket.md || https://vueuse.org/useWebSocket
 * @param url
 * @param options
 */
export function useWebSocket<Data = any>(
  url: MaybeRefOrGetter<string | URL | undefined>,
  options: UseWebSocketOptions = {}
): UseWebSocketReturn<Data> {
  const { onConnected, onDisconnected, onError, onMessage, immediate = true, autoClose = true, protocols = [] } = options

  const data: Ref<Data | null> = ref(null)
  const status = ref<WebSocketStatus>('CLOSED')
  const wsRef = ref<WebSocket | undefined>()
  const urlRef = toRef(url)

  // 心跳暂停
  let heartbeatPause: Fn | undefined
  // 心跳恢复
  let heartbeatResume: Fn | undefined

  // 连接是否已完全关闭
  let explicitlyClosed = false
  // 重连次数
  let retried = 0

  // 缓冲 data
  let bufferedData: (string | ArrayBuffer | Blob)[] = []
  // ping pong timeout 时长
  let pongTimeoutWait: ReturnType<typeof setTimeout> | undefined

  /**
   * 发送缓冲数据
   */
  const _sendBuffer = () => {
    if (bufferedData.length && wsRef.value && status.value === 'OPEN') {
      for (const buffer of bufferedData) {
        wsRef.value.send(buffer)
      }
      // 清空缓冲
      bufferedData = []
    }
  }

  /**
   * 重置心跳
   */
  const resetHeartbeat = () => {
    clearTimeout(pongTimeoutWait)
    pongTimeoutWait = undefined
  }

  // Status code 1000 -> Normal closure https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code
  /**
   * 正常关闭
   * @param code
   * @param reason
   */
  const close: WebSocket['close'] = (code = 1000, reason) => {
    if (!isClient || !wsRef.value) return
    explicitlyClosed = true
    resetHeartbeat()
    heartbeatPause?.()
    wsRef.value?.close(code, reason)
  }

  /**
   * 发送数据
   * @param data
   * @param useBuffer
   */
  const send = (data: string | ArrayBuffer | Blob, useBuffer = true) => {
    if (!wsRef.value || status.value !== 'OPEN') {
      useBuffer && bufferedData.push(data)
      return false
    }
    // 先把断开连接时未发送的缓存数据先发送出去
    _sendBuffer()
    // 再发送当前数据
    wsRef.value.send(data)
    return true
  }

  /**
   * 初始化 websocket
   */
  const _init = () => {
    // 正常关闭时不重连
    if (explicitlyClosed || typeof urlRef.value === 'undefined') return

    const ws = new WebSocket(urlRef.value, protocols)
    wsRef.value = ws
    status.value = 'CONNECTING'

    /**
     * 连接成功
     */
    ws.onopen = () => {
      status.value = 'OPEN'
      onConnected?.(ws)
      heartbeatResume?.()
      // 将缓存的数据发送出去
      _sendBuffer()
    }

    /**
     * 连接断开
     * @param ev
     */
    ws.onclose = ev => {
      status.value = 'CLOSED'
      wsRef.value = undefined
      onDisconnected?.(ws, ev)

      // 不是正常关闭时，重连
      if (!explicitlyClosed && options.autoReconnect) {
        const { retries = -1, delay = 1000, onFailed } = resolveNestedOptions(options.autoReconnect)
        retried += 1

        if (typeof retries === 'number' && (retries < 0 || retried < retries)) {
          setTimeout(_init, delay)
        } else if (typeof retries === 'function' && retries()) {
          setTimeout(_init, delay)
        } else {
          onFailed?.()
        }
      }
    }

    /**
     * 连接错误
     * @param e
     */
    ws.onerror = e => {
      onError?.(ws!, e)
    }

    /**
     * 接收消息
     * @param e
     */
    ws.onmessage = (e: MessageEvent) => {
      if (options.heartbeat) {
        // 重置心跳计时
        // 以保持心跳的频率和稳定性
        resetHeartbeat()
        const { message = DEFAULT_PING_MESSAGE } = resolveNestedOptions(options.heartbeat)
        // 服务器发送的是心跳消息，return
        // 避免将心跳消息当作普通消息处理
        if (e.data === message) {
          return
        }
      }

      data.value = e.data
      onMessage?.(ws!, e)
    }
  }
  // 心跳
  if (options.heartbeat) {
    const { message = DEFAULT_PING_MESSAGE, interval = 1000, pongTimeout = 1000 } = resolveNestedOptions(options.heartbeat)

    const { pause, resume } = useIntervalFn(
      () => {
        // 发送心跳
        send(message, false)
        if (!isNullAndUnDef(pongTimeoutWait)) return
        // 设置心跳超时
        pongTimeoutWait = setTimeout(() => {
          // 超时后关闭连接
          // auto-reconnect will be trigger with ws.onclose()
          close()
          explicitlyClosed = false
        }, pongTimeout)
      },
      interval,
      { immediate: false }
    )
    heartbeatPause = pause
    heartbeatResume = resume
  }

  if (autoClose) {
    if (isClient) {
      useEventListener('beforeunload', () => close())
    }
    tryOnScopeDispose(() => close())
  }

  /**
   * 开启连接
   */
  const open = () => {
    if (!isClient && !isWorker) return

    close()
    explicitlyClosed = false
    retried = 0
    _init()
  }

  if (immediate) {
    watch(urlRef, open, { immediate: true })
  }

  return {
    data,
    status,
    close,
    send,
    open,
    ws: wsRef
  }
}
